# 设置文件夹路径   
folder_path = "pictures" 

import os  
import random  
from PIL import Image  
import matplotlib.pyplot as plt
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, default = 0)
parser.add_argument('--team', type=int, default = 1)
args = parser.parse_args()
random.seed(args.seed)

# 获取文件夹中所有图片文件的列表
image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png'))]

################################################################################################

# 加权
# image_files.append("Aatrox.png")

################################################################################################

# 随机选择24张图片  
random_image_files = random.sample(image_files, 24)

# 处理重复
set_heros=set(random_image_files) # 将原来的列表转换为集合，这样可以不计重复
while before:=len(set_heros)!=24: # 没抽满24个英雄
    sample_=random.sample(image_files,1)
    set_heros.update(sample_) # 再抽1个，添加进去
    after=len(set_heros)
    if after>before:
        random_image_files.append(sample_[0])
set_heros=sorted(set_heros, key=random_image_files.index)

random_image_files=list(set_heros)# 转回列表

if args.team==1:
    random_image_files_team =random_image_files[0:12]
else:
    random_image_files_team =random_image_files[12:24]

# 设置显示图片的网格大小 
rows, cols = 2, 6  # 2行6列，共12张图片

# 显示图片
fig, axes = plt.subplots(rows, cols, figsize=(cols*2, rows*2))  # 调整figsize以适应图片大小  
for ax, img_file in zip(axes.flat, random_image_files_team):  
    img_path = os.path.join(folder_path, img_file)  
    img = Image.open(img_path)  
    ax.imshow(img)  
    ax.axis('off')  # 关闭坐标轴  

# 调整子图之间的间距 
plt.subplots_adjust(wspace=0.1, hspace=0.1)  

# 显示整个图像
plt.show()